#!/usr/bin/env python

# Status:
   # Works as just interface
# 1/12/12 TODO:
   # Also, better object estimation if object bounding box is adjusted by what is found in the world.
   # send off world frame to dan's code -> PointStamped in wide_optical_stereo
   # Better update of rectangle after head moved

import roslib; roslib.load_manifest('brain_stuff')
import rospy
import smach
import image_geometry
import message_filters
import actionlib
from sensor_msgs.msg import Image, CameraInfo
from cv_bridge import CvBridge, CvBridgeError
from geometry_msgs.msg import PointStamped
from pr2_controllers_msgs.msg import *
from actionlib_msgs.msg import GoalStatus
import tf
import cv
import math
import copy

class MoveHead(smach.State):
   def __init__(self):
      smach.State.__init__(self, outcomes=['done', 'fail'], input_keys=['point'])
      rospy.loginfo("Waiting on head trajectory client")
      self.client = actionlib.SimpleActionClient('/head_traj_controller/point_head_action', PointHeadAction)
      self.client.wait_for_server()
      self.g = PointHeadGoal()

   def execute(self, userdata):
      rospy.loginfo("MoveHead execution")
      if userdata.point == None:
         rospy.loginfo("Something weird happened, and goal point is null. Fix me?")
         return 'fail'
 
      # Look at object
      self.g.target = userdata.point
      self.g.min_duration = rospy.Duration(1.0)
      self.client.send_goal(self.g)
      self.client.wait_for_result()
      if not self.client.get_state() == GoalStatus.SUCCEEDED:
         rospy.loginfo("INFIDEL! I REFUSE TO LOOK WHERE YOU TELL ME TO!")
         return 'fail'

      return 'done'


class Interface(smach.State):
   def __init__(self):
      smach.State.__init__(self, outcomes=['success', 'fail', 'aborted'], output_keys=['point'])
      # Flags and global vars
      self.done = False
      self.cv_image = None
      self.image_header = None
      self.active = False
      self.cam_model = None
      self.point = PointStamped()
      self.point.header = None

      # OpenCv setup
      self.window_name = "Image Window"
      cv.NamedWindow(self.window_name, 1)
      cv.SetMouseCallback(self.window_name, self.mouse_callback) 
      self.bridge = CvBridge()

      # Click bounding box
      self.center = None
      
      # Tf setup
      self.global_frame = 'base_link'
      self.listener = tf.TransformListener()

      # Publishers and Subscribers
      self.image_sub = message_filters.Subscriber("/wide_stereo/right/image_rect_color", Image)
      self.camera_info_sub = message_filters.Subscriber("/wide_stereo/right/camera_info", CameraInfo)
      self.sync = message_filters.TimeSynchronizer([self.image_sub, self.camera_info_sub], 10)
      self.sync.registerCallback(lambda x,y: self.callback(x, y))

   def execute(self, userdata):
      userdata.point = None
      self.center = None
      self.done = False
      self.active = True
      while(not self.done and not rospy.is_shutdown()):
         rospy.sleep(0.01)
      if rospy.is_shutdown():
         self.done = True
         self.active = False
         rospy.loginfo("Keyboard Interrupt received. Aborting")
         cv.DestroyAllWindows()
         return 'aborted'

      #self.center = [320,240]

      self.active = False
      userdata.point = copy.deepcopy(self.point)
      return 'success' 

   def showImage(self):
      if self.cv_image == None:
         rospy.logerror("Trying to display an image before it exists! Skipping.")
         return
      if not self.center == None:
         off = 50 #This offset just determines the size of the bounding box. In pixels
         color_rgb = (0, 255, 0)
         cv.Rectangle(self.cv_image, (self.center[0]-off, self.center[1]-off), (self.center[0]+off, self.center[1]+off), color_rgb, 1)
   
      cv.ShowImage(self.window_name, self.cv_image)
      cv.WaitKey(2) #Any less than 2 and it freezes when clicked... :(
      
   def callback(self, data, cam_info):
      if rospy.is_shutdown():
         return
      self.point.header = copy.deepcopy(data.header)
      if self.cam_model == None:
         self.cam_model = image_geometry.PinholeCameraModel()
         self.cam_model.fromCameraInfo(cam_info)
      try:
         self.cv_image = self.bridge.imgmsg_to_cv(data, "bgr8")
      except CvBridgeError, e:
         print e
         return
      self.showImage()
      
   def mouse_callback(self, event, x, y, flags, param):
      if self.done or not self.active:
         return
      if event == cv.CV_EVENT_LBUTTONDOWN:
         rospy.loginfo("Mouse button clicked. Selected ({0},{1}) in image coordinates".format(x, y))
         self.center = [x,y]
         if self.cam_model == None:
            rospy.loginfo("Something went wrong and camera model has not been initialized. Ignoring mouse feedback")
            return
         ray = self.cam_model.projectPixelTo3dRay((x,y))
         self.point.point.x = ray[0]
         self.point.point.y = ray[1]
         self.point.point.z = ray[2]
         self.done = True
         

if __name__== "__main__":
   rospy.init_node('image_converter', log_level=rospy.DEBUG)

   machine = smach.StateMachine(outcomes=['success', 'fail', 'aborted'])
   with machine:
      smach.StateMachine.add('Display Image', Interface(), transitions={'success':'Move Head', 'fail':'fail', 'aborted':'aborted'})
      smach.StateMachine.add('Move Head', MoveHead(), transitions={'done':'Display Image', 'fail':'Display Image'})
   try:
      rospy.loginfo("Starting state machine execution")
      machine.execute()
      rospy.spin()
   except KeyboardInterrupt:
      rospy.loginfo("Shutting Down")
   cv.DestroyAllWindows()


